# -*- coding: utf-8 -*-
import re
import time
import json
import traceback
import uuid
from typing import List, Dict, Union, Generator

from loguru import logger
from openai import OpenAI

from db.repository.clause import db_save_compare_results, db_get_standard_by_clauseId
from utils.redis_utils import redis_pool


def model_chat(compare_prompt: str, stream: bool = True, json: bool = True):
    # api_key = "4130db91-11f6-44cd-b3bb-71584e47c61d"
    # base_url = "https://ark.cn-beijing.volces.com/api/v3"
    # llm_id = "ep-20250226184914-c495r"
    api_key = "sk-6f89326783dd42b3b9c63284ef454cf2"
    base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    llm_id = "qwen3-32b"
    client = OpenAI(api_key=api_key, base_url=base_url)

    messages = [
        {"role": "user", "content": compare_prompt}
    ]
    if stream:
        response = client.chat.completions.create(
            model=llm_id,
            messages=messages,
            temperature=0.7,
            stream=True,
            extra_body={"enable_thinking": False}
        )
        for chunk in response:
            chunk_message = chunk.choices[0].delta.content
            if chunk_message:
                yield chunk_message
    else:
        response = client.chat.completions.create(
            model=llm_id,
            messages=messages,
            temperature=0.7,
            stream=False,
            extra_body={"enable_thinking": True}
        )
        content = response.choices[0].message.content
        yield content


CLAUSE_CMP_PROMPT_V2 = """你是港口安全管理专家，现要求你做企业制度的合规性审核，你将结合自己的专业知识和能力，给出最权威的审核意见。
标准条款内容如下：{standard_clause}
企业制度为：{enterprise_clause}
比对与输出要求：
1. 依次(这里的依次指的是法规条款的顺序,必须一个不落)参考标准条款中的内容(需要参考的内容为："法规条款"、"人工优化过的内容")，从企业制度中找到与之相关的待审核条款(指明条目)，进行比对。
2. 比对时注意：
	a. 企业制度是否违背标准条款精神,法规要素、主题、限定范围、数值等是否合理；
	b. 企业制度的限定、要求比标准条款严格是没问题的，宽松的话需要指出；
	c. 最终返回检查结果为一个json列表, 列表中为逐一比对的结果，按照法规条款的顺序输出,必须一个不落！有20条就得输出20条，有30条就得输出30条。e.g. [{"id": "id(法规条款的id)", "regulation":"法规制度(原封不动输出)", "enterprise_clause":"企业制度"（指明条目序号(取文章结构的最高两级,如，第三章第二条、第2点下的2.3)，若有多个，则合并输出), "result":"检查结果", "advice":"修改建议"}... ];
	d. 检查结果有三种：符合、不符合、内容缺失(企业制度中没有相关内容)，如果符合，修改建议为：无。
在基于标准制度定位企业制度的时候，千万要准，不能出错，后续别人会基于你的定位进一步处理！
json列表为:
"""


def extract_json_from_stream(stream_generator) -> Union[Dict, List]:
    """
    从流式输出中提取完整的JSON对象或数组

    Args:
        stream_generator: 生成文本块的生成器

    Yields:
        解析出的JSON对象(dict)或数组(list)
    """
    buffer = ""
    stack = []
    in_json = False
    json_start_char = None  # 记录JSON起始字符是 '{' 还是 '['

    for chunk in stream_generator:
        buffer += chunk

        for char in chunk:
            if char in ("{", "["):
                if not stack:  # 如果是第一个括号，记录类型
                    json_start_char = char
                stack.append(char)
                in_json = True
            elif char in ("}", "]"):
                if stack:
                    stack.pop()
                else:
                    # 多余的右括号，重置状态
                    buffer = ""
                    in_json = False

        # 当栈为空且处于JSON模式时尝试解析
        if in_json and not stack and json_start_char:
            try:
                # 查找最外层的括号对
                start = buffer.find(json_start_char)
                end = buffer.rfind("}" if json_start_char == "{" else "]")

                if start != -1 and end != -1 and end > start:
                    json_str = buffer[start:end + 1]
                    try:
                        data = json.loads(json_str)
                        yield data
                        buffer = buffer[end + 1:]  # 保留未处理的部分
                        in_json = False
                        json_start_char = None
                    except json.JSONDecodeError:
                        # JSON不完整，继续累积数据
                        continue
            except Exception as e:
                logger.error(f"JSON解析错误: {str(e)}")
                buffer = ""
                in_json = False
                json_start_char = None

    # 处理最后可能剩余的JSON
    if in_json and buffer and json_start_char:
        try:
            data = json.loads(buffer)
            yield data
        except json.JSONDecodeError:
            logger.error(f"最终JSON解析失败: {buffer}")


def compare_clauses_with_standard(standard_clauses: List, parsed_docx: List[str], file_obj, user_id):
    """
    str(file_id): {"total": 20, "processed": 10, "parsed_docx": [str,...], "results": [str,...]}
    """
    file_id = file_obj.id

    start = time.time()
    try:
        redis_pool.hset(str(file_id), "total", len(standard_clauses))
        redis_pool.hset(str(file_id), "processed", 1)
        redis_pool.hset(str(file_id), "parsed_docx", json.dumps(parsed_docx, ensure_ascii=False))
        compare_prompt = CLAUSE_CMP_PROMPT_V2.replace("{standard_clause}", str(standard_clauses)).replace("{enterprise_clause}", "".join(parsed_docx))

        json_list = []
        processed = 0
        for json_obj in extract_json_from_stream(model_chat(compare_prompt)):
            logger.info("解析到有效JSON:")
            logger.info(json_obj)
            if isinstance(json_obj, Dict):
                json_list.append(json_obj)
                processed += 1
            if isinstance(json_obj, List):
                json_list.extend(json_obj)
                processed += len(json_obj)

            redis_pool.hset(str(file_id), "processed", processed)

        redis_pool.hset(str(file_id), "results", json.dumps(json_list, ensure_ascii=False))
    except Exception as e:
        logger.error(e)


EXTRACT_CONTENT_PROMPT = """
我会给你一个完整的企业制度文档和具体的条目序号，请你帮我从文档中准确找出对应条目的具体内容。

企业制度文档：
{enterprise_doc}

需要查找的条目序号：{clause_number}

请只返回对应条目的具体内容，不需要任何解释或评论。如果找不到完全匹配的条目，返回空字符串。
最终结果返回json列表，格式为[{"条目序号":"", "具体内容原文": ""},...]
"""


def conclusion_gen(con_li: List):
    c = 0
    query_in = ''
    for con in con_li:
        c += 1
        institution = con.get("institution", "")
        clause_name = con.get("clause_name", "")
        human_deal_text = con.get("human_deal_text", "")
        institution = con.get("institution", "")
        related_chunks = con.get("related_chunks", {})
        # item = str(c) + '. **' + institution + '**:' + '\n' + '   - ' + clause_name + '规定了' + human_deal_text + '\n'
        item = f"{str(c)}.**{institution}**:\n   - {clause_name}规定了{human_deal_text}\n"
        if con['review_result'] == '存在缺失':
            add = '   - 上传的制度缺少对应条目，建议补充'
        else:
            add = '   - 上传的制度提到了' + related_chunks.get("content", "") + '与标准制度对应条目有矛盾' + con[
                'advice'] + '\n'
        item += add
        query_in += item
    conclusion_gen_prompt = """对如下的制度审核小结进行润色，要求：
a. 描述流畅简洁，描述制度内容的时候，精简差异之外的描述，说明要点，强调标准制度条目与上传制度条目的差异及修改建议；
b. 直接输出润色后的内容；
c. 示例：
	1. **安全培训的组织实施**：
	   - 制度A规定了安全培训的组织实施，包括资金保障、培训资质、培训计划的制定和审核、培训的执行和评估等。
	   - 上传的制度提到了年度安全培训计划的制定和审批，但未详细说明组织实施的具体流程和要求，建议添加组织实施的具体流程和要求。
	2. **安全培训档案的建立和管理**：
	   - 制度A提到了安全培训档案的建立和管理，要求如实、详细、准确记录培训考核情况，并进行月度动态管理。
	   - 上传的制度中虽然提到了培训档案的建立，但未提及月度动态管理的要求，建议添加月度动态管理的要求。
d. 小结内容为{conclusion}
直接输出润色后的内容为："""
    tmp_res = model_chat(conclusion_gen_prompt.replace("{conclusion}", query_in))
    res = ''
    for r in tmp_res:
        res += r
    return res


async def iter_chat(clause_file, comparison_results, enterprise_doc, user_id):
    final_results = []
    name = clause_file.origin_name.rsplit('.', 1)[0]

    try:
        his_id = uuid.uuid1()
        yield json.dumps({"id": str(his_id)})

        for item in comparison_results:
            enterprise_clause = item.get('enterprise_clause', '无')
            enterprise_clause_id = item.get("id")
            standard_file_obj = db_get_standard_by_clauseId(enterprise_clause_id)

            if not standard_file_obj:
                continue

            output_json = {
                "id": str(standard_file_obj.id),
                "institution": standard_file_obj.institution,
                "clause_name": standard_file_obj.clause_name,
                "clause_number": item.get("regulation", ""),
                "original_text": standard_file_obj.original_text,
                "review_result": item.get("result", ""),
                "advice": item.get("advice", ""),
                "human_deal_text": standard_file_obj.artificial_text
            }

            if enterprise_clause == "无" or output_json.get("result") == "内容缺失":
                extracted_content = "存在缺失"
                output_json["related_chunks"] = {"file_id": str(clause_file.id), "file_name": name, "content": ""}  # noqa
                output_json["review_result"] = extracted_content
                output_json['advice'] = '请添加相关要素条款'
                logger.info(output_json)
                yield json.dumps(output_json, ensure_ascii=False)
            else:
                compare_prompt = EXTRACT_CONTENT_PROMPT.replace(
                    "{enterprise_doc}", "".join(enterprise_doc)
                ).replace(
                    "{clause_number}", enterprise_clause
                )

                response = ""
                for chunk in model_chat(compare_prompt):
                    response += chunk

                chunks = []
                try:
                    res = json.loads(response.replace("```json", "").replace("```", ""))
                    for r in res:
                        content_origin: str = r.get("具体内容原文", "").strip()
                        content_num = r.get("条目序号", "").strip()
                        if content_origin.startswith(content_num):
                            content_origin = content_origin.replace(content_num, "")
                        chunks.append({"title": content_num, "content": content_origin})
                except json.decoder.JSONDecodeError as e:
                    logger.error(f"JSONDecodeError: <Response: {response}>")

                if len(chunks):
                    output_json["related_chunks"] = {
                        "file_id": str(clause_file.id),
                        "file_name": name,
                        "contents": chunks
                    }
                    logger.info(output_json)
                    yield json.dumps(output_json, ensure_ascii=False)

            final_results.append(output_json)

        if len(final_results):
            conclusion = conclusion_gen(final_results)
        else:
            conclusion = "经port-ai制度审核助手分析，上传的制度符合标准制度规定，port-ai也可能犯错，重要信息还请进一步核查。"

        yield json.dumps({"summarize": conclusion}, ensure_ascii=False)
        final_results.append({"summarize": conclusion})
        redis_pool.delete(str(clause_file.id))

        history = db_save_compare_results(
            his_id=his_id,
            clause_file=clause_file,
            user_id=user_id,
            parsed_docx=comparison_results,
            compare_result=final_results,
        )
        logger.info(f"finished")

    except Exception as e:
        logger.error(str(traceback.format_exc()))
        yield json.dumps({"error": str(e)}, ensure_ascii=False)
